#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
去水印网络可视化训练脚本
包含实时图像显示、loss/acc曲线图等功能
"""

import os
import sys
import time
import yaml
import json
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from easydict import EasyDict
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import lpips

# 导入SepMark相关模块
from watermark_removel import WatermarkRemovalNet
from network.Dual_Mark import Network as SepMarkNetwork
from utils.img_loader import ImageFolder


class VisualTrainer:
    """带可视化功能的去水印网络训练器"""
    
    def __init__(self, config_path='cfg/train_watermark_removal.yaml'):
        self.load_config(config_path)
        self.setup_device()
        self.setup_models()
        self.setup_optimizer()
        self.setup_loss_functions()
        self.setup_dataloader()
        self.setup_visualization()
        self.setup_logging()
        
    def load_config(self, config_path):
        """加载配置文件"""
        with open(config_path, 'r') as f:
            self.config = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))
        print(f"配置加载完成: {config_path}")
        
    def setup_device(self):
        """设置设备"""
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")
        
    def setup_models(self):
        """设置模型"""
        # 去水印网络（待训练）
        self.removal_net = WatermarkRemovalNet().to(self.device)
        print(f"去水印网络参数量: {self.removal_net.count_parameters():,}")
        
        # SepMark网络（预训练，用于评估）
        self.sepmark_net = SepMarkNetwork(
            message_length=self.config.message_length,
            noise_layers_R=self.config.noise_layers_R,
            noise_layers_F=self.config.noise_layers_F,
            attention_encoder=self.config.attention_encoder,
            attention_decoder=self.config.attention_decoder,
            device=self.device
        ).to(self.device)
        
        # 加载预训练的SepMark模型
        sepmark_model_path = f"results/{self.config.sepmark_model_folder}/models/EC_{self.config.sepmark_epoch}.pth"
        if os.path.exists(sepmark_model_path):
            self.sepmark_net.load_state_dict(torch.load(sepmark_model_path, map_location=self.device))
            self.sepmark_net.eval()
            print(f"SepMark模型加载成功: {sepmark_model_path}")
        else:
            print(f"警告: SepMark模型文件不存在: {sepmark_model_path}")
            
    def setup_optimizer(self):
        """设置优化器"""
        self.optimizer = optim.Adam(
            self.removal_net.parameters(),
            lr=self.config.lr,
            betas=(self.config.beta1, 0.999)
        )
        
        # 学习率调度器
        self.scheduler = optim.lr_scheduler.StepLR(
            self.optimizer, 
            step_size=self.config.lr_decay_step, 
            gamma=self.config.lr_decay_gamma
        )
        
    def setup_loss_functions(self):
        """设置损失函数"""
        self.mse_loss = nn.MSELoss()
        self.l1_loss = nn.L1Loss()
        
        # LPIPS感知损失
        self.lpips_loss = lpips.LPIPS(net='alex').to(self.device)
        
    def setup_dataloader(self):
        """设置数据加载器"""
        # 图像变换
        self.transform = transforms.Compose([
            transforms.Resize((self.config.image_size, self.config.image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])
        
        # 数据集
        dataset = ImageFolder(self.config.dataset_path, transform=self.transform)
        self.dataloader = DataLoader(
            dataset, 
            batch_size=self.config.batch_size,
            shuffle=True, 
            num_workers=4,
            pin_memory=True
        )
        
        print(f"数据集大小: {len(dataset)}")
        
    def setup_visualization(self):
        """设置可视化"""
        plt.ion()  # 开启交互模式
        self.fig = plt.figure(figsize=(20, 12))
        self.gs = GridSpec(3, 4, figure=self.fig, hspace=0.3, wspace=0.3)
        
        # 图像显示子图
        self.ax_original = self.fig.add_subplot(self.gs[0, 0])
        self.ax_watermarked = self.fig.add_subplot(self.gs[0, 1])
        self.ax_removed = self.fig.add_subplot(self.gs[0, 2])
        self.ax_diff = self.fig.add_subplot(self.gs[0, 3])
        
        # Loss曲线子图
        self.ax_loss = self.fig.add_subplot(self.gs[1, :2])
        self.ax_acc = self.fig.add_subplot(self.gs[1, 2:])
        
        # 水印解码结果子图
        self.ax_decode_c = self.fig.add_subplot(self.gs[2, 0])
        self.ax_decode_rf = self.fig.add_subplot(self.gs[2, 1])
        self.ax_lpips = self.fig.add_subplot(self.gs[2, 2])
        self.ax_metrics = self.fig.add_subplot(self.gs[2, 3])
        
        # 设置标题
        self.ax_original.set_title('Original Image')
        self.ax_watermarked.set_title('Watermarked Image')
        self.ax_removed.set_title('Watermark Removed')
        self.ax_diff.set_title('Difference')
        self.ax_loss.set_title('Training Losses')
        self.ax_acc.set_title('Watermark Detection Accuracy')
        self.ax_decode_c.set_title('Decoder C Message')
        self.ax_decode_rf.set_title('Decoder RF Message')
        self.ax_lpips.set_title('LPIPS Score Trend')
        self.ax_metrics.set_title('Training Metrics')
        
        # 初始化数据记录
        self.loss_history = {
            'reconstruction': [],
            'lpips': [],
            'total': []
        }
        self.acc_history = {
            'decoder_c': [],
            'decoder_rf': []
        }
        self.lpips_history = []
        self.epoch_history = []
        
    def setup_logging(self):
        """设置日志记录"""
        # 创建结果文件夹
        timestamp = time.strftime("%Y_%m_%d_%H_%M_%S")
        self.result_folder = f"results/watermark_removal_{timestamp}"
        os.makedirs(self.result_folder, exist_ok=True)
        os.makedirs(f"{self.result_folder}/models", exist_ok=True)
        os.makedirs(f"{self.result_folder}/images", exist_ok=True)
        os.makedirs(f"{self.result_folder}/plots", exist_ok=True)
        
        # 保存配置文件
        with open(f"{self.result_folder}/config.yaml", 'w') as f:
            yaml.dump(dict(self.config), f)
            
        print(f"结果保存路径: {self.result_folder}")
        
    def generate_watermark_message(self, batch_size):
        """生成水印消息"""
        message = torch.Tensor(
            np.random.choice(
                [-self.config.message_range, self.config.message_range],
                (batch_size, self.config.message_length)
            )
        ).to(self.device)
        return message
        
    def embed_watermark(self, images):
        """使用SepMark嵌入水印"""
        with torch.no_grad():
            messages = self.generate_watermark_message(images.size(0))
            watermarked_images = self.sepmark_net.encoder_decoder.module.encoder(images, messages)
            return watermarked_images, messages
            
    def decode_watermark(self, images):
        """解码水印"""
        with torch.no_grad():
            decoded_c = self.sepmark_net.encoder_decoder.module.decoder_C(images)
            decoded_rf = self.sepmark_net.encoder_decoder.module.decoder_RF(images)
            return decoded_c, decoded_rf
            
    def calculate_accuracy(self, original_message, decoded_message):
        """计算水印检测准确率"""
        # 转换为二进制
        orig_binary = (original_message > 0).float()
        decoded_binary = (decoded_message > 0).float()
        
        # 计算准确率
        correct = (orig_binary == decoded_binary).float().mean()
        return correct.item()
        
    def denormalize_image(self, tensor):
        """反归一化图像用于显示"""
        return (tensor * 0.5 + 0.5).clamp(0, 1)
        
    def update_visualization(self, epoch, batch_idx, batch_data):
        """更新可视化显示"""
        original_images, watermarked_images, removed_images, messages = batch_data
        
        # 选择第一张图像用于显示
        idx = 0
        original = self.denormalize_image(original_images[idx].cpu())
        watermarked = self.denormalize_image(watermarked_images[idx].cpu())
        removed = self.denormalize_image(removed_images[idx].cpu())
        
        # 计算差异图像
        diff = torch.abs(original - removed)
        
        # 显示图像
        self.ax_original.clear()
        self.ax_watermarked.clear()
        self.ax_removed.clear()
        self.ax_diff.clear()
        
        self.ax_original.imshow(original.permute(1, 2, 0))
        self.ax_original.set_title('Original Image')
        self.ax_original.axis('off')
        
        self.ax_watermarked.imshow(watermarked.permute(1, 2, 0))
        self.ax_watermarked.set_title('Watermarked Image')
        self.ax_watermarked.axis('off')
        
        self.ax_removed.imshow(removed.permute(1, 2, 0))
        self.ax_removed.set_title('Watermark Removed')
        self.ax_removed.axis('off')
        
        self.ax_diff.imshow(diff.permute(1, 2, 0))
        self.ax_diff.set_title(f'Difference (MSE: {torch.mean(diff**2):.4f})')
        self.ax_diff.axis('off')
        
        # 解码水印并显示
        decoded_c, decoded_rf = self.decode_watermark(removed_images[:1])
        
        # 显示解码结果
        self.ax_decode_c.clear()
        self.ax_decode_rf.clear()
        
        # 计算准确率
        acc_c = self.calculate_accuracy(messages[:1], decoded_c)
        acc_rf = self.calculate_accuracy(messages[:1], decoded_rf)
        
        # 显示消息为条形图
        message_orig = messages[0].cpu().numpy()
        decoded_c_np = decoded_c[0].cpu().numpy()
        decoded_rf_np = decoded_rf[0].cpu().numpy()
        
        x_pos = range(min(32, len(message_orig)))  # 只显示前32位
        
        self.ax_decode_c.bar(x_pos, message_orig[:32], alpha=0.7, label='Original', color='blue')
        self.ax_decode_c.bar(x_pos, decoded_c_np[:32], alpha=0.7, label='Decoded', color='red')
        self.ax_decode_c.set_title(f'Decoder C (Acc: {acc_c:.3f})')
        self.ax_decode_c.legend()
        
        self.ax_decode_rf.bar(x_pos, message_orig[:32], alpha=0.7, label='Original', color='blue')
        self.ax_decode_rf.bar(x_pos, decoded_rf_np[:32], alpha=0.7, label='Decoded', color='red')
        self.ax_decode_rf.set_title(f'Decoder RF (Acc: {acc_rf:.3f})')
        self.ax_decode_rf.legend()
        
        # 更新loss曲线
        if len(self.loss_history['total']) > 0:
            self.ax_loss.clear()
            epochs = range(len(self.loss_history['total']))
            self.ax_loss.plot(epochs, self.loss_history['reconstruction'], label='Reconstruction', alpha=0.8)
            self.ax_loss.plot(epochs, self.loss_history['lpips'], label='LPIPS', alpha=0.8)
            self.ax_loss.plot(epochs, self.loss_history['total'], label='Total', alpha=0.8)
            self.ax_loss.set_xlabel('Batch')
            self.ax_loss.set_ylabel('Loss')
            self.ax_loss.set_title('Training Losses')
            self.ax_loss.legend()
            self.ax_loss.grid(True, alpha=0.3)
        
        # 更新准确率曲线
        if len(self.acc_history['decoder_c']) > 0:
            self.ax_acc.clear()
            epochs = range(len(self.acc_history['decoder_c']))
            self.ax_acc.plot(epochs, self.acc_history['decoder_c'], label='Decoder C', alpha=0.8)
            self.ax_acc.plot(epochs, self.acc_history['decoder_rf'], label='Decoder RF', alpha=0.8)
            self.ax_acc.set_xlabel('Batch')
            self.ax_acc.set_ylabel('Accuracy')
            self.ax_acc.set_title('Watermark Detection Accuracy')
            self.ax_acc.legend()
            self.ax_acc.grid(True, alpha=0.3)
            
        # 更新LPIPS曲线
        if len(self.lpips_history) > 0:
            self.ax_lpips.clear()
            epochs = range(len(self.lpips_history))
            self.ax_lpips.plot(epochs, self.lpips_history, label='LPIPS Score', alpha=0.8, color='green')
            self.ax_lpips.set_xlabel('Batch')
            self.ax_lpips.set_ylabel('LPIPS')
            self.ax_lpips.set_title('LPIPS Score Trend')
            self.ax_lpips.legend()
            self.ax_lpips.grid(True, alpha=0.3)
            
        # 显示训练指标
        self.ax_metrics.clear()
        current_lr = self.optimizer.param_groups[0]['lr']
        metrics_text = f"""Epoch: {epoch}
Batch: {batch_idx}
Learning Rate: {current_lr:.6f}
Decoder C Acc: {acc_c:.3f}
Decoder RF Acc: {acc_rf:.3f}
Total Parameters: {self.removal_net.count_parameters():,}"""
        
        self.ax_metrics.text(0.1, 0.5, metrics_text, transform=self.ax_metrics.transAxes, 
                           fontsize=12, verticalalignment='center')
        self.ax_metrics.axis('off')
        
        plt.draw()
        plt.pause(0.001)
        
    def save_checkpoint(self, epoch, loss):
        """保存模型检查点"""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.removal_net.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'loss': loss,
            'loss_history': self.loss_history,
            'acc_history': self.acc_history,
            'lpips_history': self.lpips_history
        }
        
        torch.save(checkpoint, f"{self.result_folder}/models/removal_net_epoch_{epoch}.pth")
        
        # 保存最佳模型
        if not hasattr(self, 'best_loss') or loss < self.best_loss:
            self.best_loss = loss
            torch.save(checkpoint, f"{self.result_folder}/models/removal_net_best.pth")
            
    def save_plots(self, epoch):
        """保存图表"""
        plt.savefig(f"{self.result_folder}/plots/training_progress_epoch_{epoch}.png", 
                   dpi=150, bbox_inches='tight')
        
    def train_epoch(self, epoch):
        """训练一个epoch"""
        self.removal_net.train()
        epoch_loss = 0.0
        epoch_acc_c = 0.0
        epoch_acc_rf = 0.0
        
        pbar = tqdm(self.dataloader, desc=f"Epoch {epoch}")
        
        for batch_idx, (images, _) in enumerate(pbar):
            images = images.to(self.device)
            
            # 生成带水印图像
            watermarked_images, messages = self.embed_watermark(images)
            
            # 去水印处理
            self.optimizer.zero_grad()
            removed_images = self.removal_net(watermarked_images)
            
            # 计算损失
            reconstruction_loss = self.l1_loss(removed_images, images)
            
            # LPIPS感知损失
            lpips_loss = self.lpips_loss(removed_images, images).mean()
            
            # 总损失
            total_loss = (self.config.lambda_reconstruction * reconstruction_loss + 
                         self.config.lambda_lpips * lpips_loss)
            
            # 反向传播
            total_loss.backward()
            self.optimizer.step()
            
            # 记录损失
            self.loss_history['reconstruction'].append(reconstruction_loss.item())
            self.loss_history['lpips'].append(lpips_loss.item())
            self.loss_history['total'].append(total_loss.item())
            self.lpips_history.append(lpips_loss.item())
            
            # 计算水印检测准确率
            decoded_c, decoded_rf = self.decode_watermark(removed_images)
            acc_c = self.calculate_accuracy(messages, decoded_c)
            acc_rf = self.calculate_accuracy(messages, decoded_rf)
            
            self.acc_history['decoder_c'].append(acc_c)
            self.acc_history['decoder_rf'].append(acc_rf)
            
            # 更新可视化（每10个batch更新一次）
            if batch_idx % 10 == 0:
                batch_data = (images, watermarked_images, removed_images, messages)
                self.update_visualization(epoch, batch_idx, batch_data)
                
            # 更新进度条
            pbar.set_postfix({
                'Loss': f'{total_loss.item():.4f}',
                'Rec': f'{reconstruction_loss.item():.4f}',
                'LPIPS': f'{lpips_loss.item():.4f}',
                'Acc_C': f'{acc_c:.3f}',
                'Acc_RF': f'{acc_rf:.3f}'
            })
            
            epoch_loss += total_loss.item()
            epoch_acc_c += acc_c
            epoch_acc_rf += acc_rf
            
        # 计算epoch平均值
        avg_loss = epoch_loss / len(self.dataloader)
        avg_acc_c = epoch_acc_c / len(self.dataloader)
        avg_acc_rf = epoch_acc_rf / len(self.dataloader)
        
        return avg_loss, avg_acc_c, avg_acc_rf
        
    def train(self):
        """主训练循环"""
        print("开始训练去水印网络...")
        print(f"目标：降低水印检测准确率，同时保持图像质量")
        
        for epoch in range(self.config.num_epochs):
            print(f"\n{'='*50}")
            print(f"Epoch {epoch+1}/{self.config.num_epochs}")
            print(f"{'='*50}")
            
            # 训练一个epoch
            avg_loss, avg_acc_c, avg_acc_rf = self.train_epoch(epoch)
            
            # 更新学习率
            self.scheduler.step()
            
            # 保存模型
            if (epoch + 1) % self.config.save_interval == 0:
                self.save_checkpoint(epoch, avg_loss)
                self.save_plots(epoch)
                
            print(f"Epoch {epoch+1} 完成:")
            print(f"  平均损失: {avg_loss:.4f}")
            print(f"  平均准确率 C: {avg_acc_c:.3f}")
            print(f"  平均准确率 RF: {avg_acc_rf:.3f}")
            print(f"  当前学习率: {self.optimizer.param_groups[0]['lr']:.6f}")
            
        print("\n训练完成！")
        print(f"结果保存在: {self.result_folder}")
        
        # 保存最终图表
        plt.savefig(f"{self.result_folder}/final_training_results.png", 
                   dpi=300, bbox_inches='tight')
        
        # 保存训练历史
        training_history = {
            'loss_history': self.loss_history,
            'acc_history': self.acc_history,
            'lpips_history': self.lpips_history,
            'config': dict(self.config)
        }
        
        with open(f"{self.result_folder}/training_history.json", 'w') as f:
            json.dump(training_history, f, indent=2)


def main():
    """主函数"""
    # 检查配置文件
    config_path = 'cfg/train_watermark_removal.yaml'
    if not os.path.exists(config_path):
        print(f"配置文件不存在: {config_path}")
        print("请先创建配置文件")
        return
        
    # 创建训练器并开始训练
    trainer = VisualTrainer(config_path)
    trainer.train()


if __name__ == "__main__":
    main() 